#![deny(warnings)]

use gateway_service::{Service, SERVICEBOOK};
use http::{header, Method, Request, Response, Version};

use std::net::SocketAddr;
use hyper::{Body};
use hyper::server::conn::Http;
use hyper::service::service_fn;
use hyper::upgrade::{OnUpgrade, Upgraded};


use tokio::net::{TcpListener, TcpStream};

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    tracing_subscriber::fmt().init();

    let in_addr: SocketAddr = ([127, 0, 0, 1], 3000).into();

    SERVICEBOOK.write().await.join_service(
        "test".into(),
        Service {
            host: "127.0.0.1".into(),
            port: 7878,
            name: "ZHRD".into(),
        },
    )?;

    let listener = TcpListener::bind(in_addr).await?;

    tracing::info!("Listening on http://127.0.0.1:3000");

    loop {
        let (stream, _) = listener.accept().await?;

        tokio::task::spawn(async move {
            if let Err(err) = Http::new()
                .http1_preserve_header_case(true)
                .http1_title_case_headers(true)
                .http2_enable_connect_protocol()
                .serve_connection(stream, service_fn(proxy))
                .with_upgrades()
                .await
            {
                println!("Failed to serve connection: {:?}", err);
            }
        });
    }
}

#[allow(dead_code)]
async fn proxy(mut req: Request<Body>) -> Result<Response<Body>, hyper::Error> {
    tracing::info!("req: {:?}", req);

    let addr = SERVICEBOOK
        .write()
        .await
        .choose_service("test".into())
        .unwrap()
        .get_url();

    if req.headers().get(header::CONNECTION) == Some(&header::HeaderValue::from(header::CONNECTION)) || Method::CONNECT == req.method() {
        tracing::info!("upgrade request");


        let stream = TcpStream::connect(addr.clone()).await.unwrap();

        let (mut sender, conn) = hyper::client::conn::handshake(stream).await?;
        tokio::task::spawn(async move {
            if let Err(err) = conn.await {
                println!("Connection failed: {:?}", err);
            }
        });
        let req_on_upgrade = req.extensions_mut()
            .remove::<OnUpgrade>().unwrap();
        let mut resp = sender.send_request(req).await?;
        let res_on_upgrade = resp.extensions_mut()
            .remove::<OnUpgrade>().unwrap();

        tokio::task::spawn(async {
            match req_on_upgrade.await {
                Ok(req_upgraded) => {
                    match res_on_upgrade.await {
                        Ok(res_upgraded) => {
                            if let Err(e) = tunnel_upgrade(req_upgraded, res_upgraded).await {
                                tracing::error!("tunnel error:{}", e);
                            }
                        }
                        Err(e) => {
                            tracing::error!("res_on_upgrade error:{}", e);
                        }
                    }
                }
                Err(e) => {
                    tracing::error!("req_on_upgrade error:{}", e);
                }
            }
        });
        Ok(resp)
    } else {
        let stream = TcpStream::connect(addr.clone()).await.unwrap();
        let (mut sender, conn) = match req.version() {
            Version::HTTP_2 => {
                tracing::info!("http2 request");
                hyper::client::conn::Builder::new().http2_only(true).handshake(stream).await?
            }
            _ => {
                tracing::info!("http request");
                hyper::client::conn::handshake(stream).await?
            }
        };

        tokio::task::spawn(async move {
            if let Err(err) = conn.await {
                println!("Connection failed: {:?}", err);
            }
        });

        sender.send_request(req).await
    }
}

async fn tunnel_upgrade(mut upgraded: Upgraded, mut res_upgraded: Upgraded) -> std::io::Result<()> {
    tracing::info!("tunnel2 connected");

    // Proxying data
    let (from_client, from_server) =
        tokio::io::copy_bidirectional(&mut upgraded, &mut res_upgraded).await?;

    // Print message when done
    tracing::info!(
        "client wrote {} bytes and received {} bytes",
        from_client, from_server,
    );

    Ok(())
}
